#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查是否为root用户
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要root权限运行"
        exit 1
    fi
}

# 检查依赖工具
check_dependencies() {
    local deps=("taskset" "pgrep" "ps")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            log_error "缺少依赖工具: $dep"
            exit 1
        fi
    done
}

# 绑定中断到NUMA0
bind_interrupts_to_numa0() {
    log_info "开始绑定中断到NUMA0..."
    
    # 获取NUMA0的CPU核心范围
    local numa0_cpus=$(get_numa_cpus 0)
    log_info "NUMA0 CPU核心范围: $numa0_cpus"
    
    # 获取virtio0相关的中断号
    local irq_numbers=$(cat /proc/interrupts 2>/dev/null | grep -E "virtio0" | awk -F ':' '{print $1}' | tr -d ' ')
    
    if [[ -z "$irq_numbers" ]]; then
        log_warning "未找到virtio0相关的中断"
        return 1
    fi
    
    local irq_count=0
    local success_count=0
    local failed_count=0
    
    # 将NUMA0的CPU范围转换为CPU列表（用于分配中断）
    local numa0_cpu_list=""
    if [[ "$numa0_cpus" =~ ^([0-9]+)-([0-9]+)$ ]]; then
        local start_cpu=${BASH_REMATCH[1]}
        local end_cpu=${BASH_REMATCH[2]}
        numa0_cpu_list=$(seq -s ',' $start_cpu $end_cpu)
    else
        numa0_cpu_list="$numa0_cpus"
    fi
    
    log_info "发现以下virtio0中断号:"
    echo "$irq_numbers" | while read -r irq_no; do
        if [[ -n "$irq_no" ]]; then
            echo "  IRQ $irq_no"
        fi
    done
    
    # 逐个绑定中断
    echo "$irq_numbers" | while read -r irq_no; do
        if [[ -n "$irq_no" ]]; then
            ((irq_count++))
            local smp_affinity_file="/proc/irq/$irq_no/smp_affinity_list"
            
            if [[ ! -f "$smp_affinity_file" ]]; then
                log_warning "中断 $irq_no 的smp_affinity_list文件不存在"
                ((failed_count++))
                continue
            fi
            
            # 读取当前的亲和性设置
            local current_affinity=$(cat "$smp_affinity_file" 2>/dev/null)
            
            # 尝试设置中断亲和性到NUMA0
            if echo "$numa0_cpu_list" > "$smp_affinity_file" 2>/dev/null; then
                local new_affinity=$(cat "$smp_affinity_file" 2>/dev/null)
                log_success "中断 $irq_no: $current_affinity -> $new_affinity"
                ((success_count++))
            else
                log_error "设置中断 $irq_no 亲和性失败"
                ((failed_count++))
            fi
        fi
    done
    
    # 使用子进程来获取最终统计，因为while循环在管道中运行在子shell中
    local total_irqs=$(echo "$irq_numbers" | grep -c .)
    local actual_success=0
    local actual_failed=0
    
    echo "$irq_numbers" | while read -r irq_no; do
        if [[ -n "$irq_no" ]]; then
            local smp_affinity_file="/proc/irq/$irq_no/smp_affinity_list"
            if [[ -f "$smp_affinity_file" ]]; then
                local current_affinity=$(cat "$smp_affinity_file" 2>/dev/null)
                # 检查是否包含NUMA0的CPU
                if [[ "$current_affinity" =~ [0-9] ]]; then
                    ((actual_success++))
                else
                    ((actual_failed++))
                fi
            else
                ((actual_failed++))
            fi
        fi
    done
    
    log_info "中断绑定完成: 总计 $total_irqs 个中断"
    
    return 0
}

# 显示中断绑定状态
show_interrupt_status() {
    log_info "显示当前中断绑定状态..."
    
    echo -e "\n${BLUE}=== 中断绑定状态 ===${NC}"
    
    # 获取virtio0相关的中断号和当前绑定状态
    local irq_numbers=$(cat /proc/interrupts 2>/dev/null | grep -E "virtio0" | awk -F ':' '{print $1}' | tr -d ' ')
    
    if [[ -z "$irq_numbers" ]]; then
        echo -e "${YELLOW}未找到virtio0相关的中断${NC}"
        return 1
    fi
    
    echo -e "\n${YELLOW}virtio0中断绑定状态:${NC}"
    
    echo "$irq_numbers" | while read -r irq_no; do
        if [[ -n "$irq_no" ]]; then
            local smp_affinity_file="/proc/irq/$irq_no/smp_affinity_list"
            if [[ -f "$smp_affinity_file" ]]; then
                local current_affinity=$(cat "$smp_affinity_file" 2>/dev/null)
                local irq_desc=$(cat /proc/interrupts 2>/dev/null | grep "^[[:space:]]*$irq_no:" | sed 's/^[[:space:]]*[0-9]*:[[:space:]]*//' | awk '{for(i=NF;i>=1;i--) if($i ~ /virtio0/) {for(j=i;j<=NF;j++) printf "%s ", $j; break}}')
                echo "  IRQ $irq_no -> CPU: $current_affinity ($irq_desc)"
            else
                echo "  IRQ $irq_no -> 无法读取绑定状态"
            fi
        fi
    done
}

bind_process_without_threads() {
    local pid=$1
    local cpu_mask=$2
    local process_name=$3
    
    if [[ ! -d "/proc/$pid" ]]; then
        log_warning "进程 $process_name (PID: $pid) 不存在"
        return 1
    fi
    
    # 绑定主进程
    if taskset -cp "$cpu_mask" "$pid" >/dev/null 2>&1; then
        log_success "绑定进程 $process_name (PID: $pid) 到CPU核心: $cpu_mask"
    else
        log_error "绑定进程 $process_name (PID: $pid) 失败"
        return 1
    fi

    return 0
}


# 将进程及其所有线程绑定到指定CPU核心
bind_process() {
    local pid=$1
    local cpu_mask=$2
    local process_name=$3
 
    if [[ ! -d "/proc/$pid" ]]; then
       log_warning "进程 $process_name (PID: $pid) 不存在"
       return 1
    fi
 
     # 绑定主进程
     if taskset -cp "$cpu_mask" "$pid" >/dev/null 2>&1; then
        log_success "绑定进程 $process_name (PID: $pid) 到CPU核心: $cpu_mask"
     else
         log_error "绑定进程 $process_name (PID: $pid) 失败"
         return 1
     fi
    
     return 0
 }
2


# 将进程及其所有线程绑定到指定CPU核心
bind_process_and_threads() {
    local pid=$1
    local cpu_mask=$2
    local process_name=$3
    
    if [[ ! -d "/proc/$pid" ]]; then
        log_warning "进程 $process_name (PID: $pid) 不存在"
        return 1
    fi
    
    # 绑定主进程
    if taskset -cap "$cpu_mask" "$pid" >/dev/null 2>&1; then
        log_success "绑定进程 $process_name (PID: $pid) 到CPU核心: $cpu_mask"
    else
        log_error "绑定进程 $process_name (PID: $pid) 失败"
        return 1
    fi
    
    # 获取所有线程ID并绑定
    : 'local thread_count=0
    if [[ -d "/proc/$pid/task" ]]; then
        for tid in /proc/$pid/task/*; do
            if [[ -d "$tid" ]]; then
                local thread_id=$(basename "$tid")
                if taskset -cp "$cpu_mask" "$thread_id" >/dev/null 2>&1; then
                    ((thread_count++))
                else
                    log_warning "绑定线程 $thread_id 失败"
                fi
            fi
        done
        log_info "  └── 绑定了 $thread_count 个线程"
    fi'
    
    return 0
}

# 根据NUMA节点获取CPU核心列表
get_numa_cpus() {
    local numa_id=$1
    local start_core=$((numa_id * 40))
    local end_core=$((start_core + 39))
    echo "$start_core-$end_core"
}


# 获取任务组的CPU核心范围
get_driver_cpu_range() {
    local task_id=$1
    task_id=$((task_id % 16)) # 0-15
    
    local base_core=$((124 + task_id * 4))  # 从124开始,driver独占1个cluser
    
    local start_core=$((base_core + 0))
    local end_core=$((base_core + 3))
    
    echo "$start_core-$end_core"
}

# 获取forwordCPU核心范围
get_forward_cpu_range() {
    local task_id=$1
    task_id=$((task_id % 16)) # 0-15
    
    local base_core=$((192 + task_id * 4))  # 192开始, 每个cluster1个forward

    local start_core=$((base_core))
    local end_core=$((base_core))
    
    echo "$start_core-$end_core"
}


# 获取other进程CPU核心范围
get_other_cpu_range() {
    
    local start_core=40
    local end_core=59
    
    echo "$start_core-$end_core"
}


# 获取workerCPU核心范围
get_worker_cpu_range() {
    local task_id=$1
    task_id=$((task_id % 16)) # 0-15
    
    local base_core=$((60 + task_id * 4))  # 从60开始,worker独占1个cluser

    local start_core=$((base_core + 0))
    local end_core=$((base_core + 3))
    
    echo "$start_core-$end_core"
}

find_highest_cpu_child_process() {
    local parent_pid=$1
    local max_cpu=0
    local highest_pid=""
    
    if [[ ! -d "/proc/$parent_pid/task" ]]; then
        return
    fi
    
    # 声明关联数组存储第一次和第二次的CPU时间
    declare -A cpu_stats1
    declare -A cpu_stats2
    declare -A valid_tids
    
    # 第一次批量读取所有线程的CPU统计，同时检查wchan
    for task_pid in /proc/$parent_pid/task/*; do
        if [[ -d "$task_pid" ]]; then
            local tid=$(basename "$task_pid")
            # 跳过父进程本身
            if [[ "$tid" != "$parent_pid" ]]; then
                # 检查wchan文件是否为0 (表示进程正在运行，未阻塞)
                local wchan=$(cat /proc/$tid/wchan 2>/dev/null)
                if [[ "$wchan" == "0" ]]; then
                    local stat=$(cat /proc/$tid/stat 2>/dev/null | awk '{print $14+$15}')
                    if [[ -n "$stat" ]]; then
                        cpu_stats1[$tid]=$stat
                        valid_tids[$tid]=1
                    fi
                fi
            fi
        fi
    done
    
    # 如果没有找到任何符合条件的子线程，直接返回
    if [[ ${#cpu_stats1[@]} -eq 0 ]]; then
        return
    fi
    
    # 短暂休眠
    sleep 0.1
    
    # 第二次批量读取所有线程的CPU统计，再次检查wchan
    for tid in "${!valid_tids[@]}"; do
        # 检查线程是否仍然存在，并且wchan仍为0
        if [[ -f "/proc/$tid/stat" && -f "/proc/$tid/wchan" ]]; then
            local wchan=$(cat /proc/$tid/wchan 2>/dev/null)
            if [[ "$wchan" == "0" ]]; then
                local stat=$(cat /proc/$tid/stat 2>/dev/null | awk '{print $14+$15}')
                if [[ -n "$stat" ]]; then
                    cpu_stats2[$tid]=$stat
                fi
            fi
        fi
    done
    
    # 计算CPU使用率差值，找到最大的
    for tid in "${!cpu_stats1[@]}"; do
        if [[ -n "${cpu_stats2[$tid]}" ]]; then
            local cpu_diff=$((cpu_stats2[$tid] - cpu_stats1[$tid]))
            if [[ "$cpu_diff" -gt "$max_cpu" ]]; then
                max_cpu=$cpu_diff
                highest_pid=$tid
            fi
        fi
    done
    
    echo "$highest_pid"
}

# 绑定worker进程
bind_worker_processes() {
    log_info "开始绑定Worker进程..."
    
    for i in {0..15}; do
        local process_pattern="xLLM::XWorker_DP_$i"
        # 使用精确匹配，避免匹配到其他数字后缀
        local pids=$(pgrep -f "xLLM::XWorker_DP_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM::XWorker_DP_${i}$")
        
        if [[ -z "$pids" ]]; then
            log_warning "未找到进程: $process_pattern"
            continue
        fi
        
        local cpu_range=$(get_worker_cpu_range "$i")
        
        echo "$pids" | while read -r pid; do
            bind_process_and_threads "$pid" "$cpu_range" "$process_pattern"
        done

        # bind high
        echo "$pids" | while read -r pid; do
            local high_cpu_worker_pid=$(find_highest_cpu_child_process "$pid")
            local high_cpu_worker_cpu_range=$(get_forward_cpu_range "$i")
            bind_process_without_threads "$high_cpu_worker_pid" "$high_cpu_worker_cpu_range" "${process_pattern}_forwrad_worker_thread"
        done
    done
}

# 绑定driver进程
bind_driver_processes() {
    log_info "开始绑定Driver进程..."
    
    for i in {0..15}; do
        local process_pattern="xLLM::Driver_$i"
        # 使用精确匹配，避免匹配到其他数字后缀
        local pids=$(pgrep -f "xLLM::Driver_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM::Driver_${i}$")
        
        if [[ -z "$pids" ]]; then
            log_warning "未找到进程: $process_pattern"
            continue
        fi
        
        local cpu_range=$(get_driver_cpu_range "$i" "driver")
        
        echo "$pids" | while read -r pid; do
            bind_process_and_threads "$pid" "$cpu_range" "$process_pattern"
        done
    done
}

# 获取collector进程pid
find_npu_collector_pid_with_top() {
    # 使用top命令的批处理模式，结合grep和awk提取PID
    local pid=$(top -b -n 1 | grep "npu_collector_s" | head -n 1 | awk '{print $1}')
    
    # 检查是否找到进程
    if [[ -n "$pid" && "$pid" =~ ^[0-9]+$ ]]; then
        echo "$pid"
        return 0
    else
        echo "未找到npu_collector_s进程"
        return 1
    fi
}

# 获取node进程pid
find_node_pid_with_top() {
    # 使用top命令的批处理模式，结合grep和awk提取PID
    local pid=$(top -b -n 1 | grep "node" | head -n 1 | awk '{print $1}')
    
    # 检查是否找到进程
    if [[ -n "$pid" && "$pid" =~ ^[0-9]+$ ]]; then
        echo "$pid"
        return 0
    else
        echo "未找到node进程"
        return 1
    fi
}



# 绑定其他进程到NUMA1
bind_other_processes() {
    log_info "开始绑定其他进程到NUMA1..."
    
    local numa1_cpus=$(get_other_cpu_range)
    
    # 绑定Servicer进程
    local servicer_pids=$(pgrep -f "xLLM::Servicer")
    if [[ -n "$servicer_pids" ]]; then
        echo "$servicer_pids" | while read -r pid; do
            bind_process_and_threads "$pid" "$numa1_cpus" "xLLM::Servicer"
        done
    else
        log_warning "未找到xLLM::Servicer进程"
    fi
    
    # 绑定RpcProxy进程
    for i in {0..15}; do
        local rpc_pattern="xLLM:RpcProxy_$i"
        # 使用精确匹配，避免匹配到其他数字后缀
        local rpc_pids=$(pgrep -f "xLLM:RpcProxy_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM:RpcProxy_${i}$")
        
        if [[ -n "$rpc_pids" ]]; then
            echo "$rpc_pids" | while read -r pid; do
                bind_process_and_threads "$pid" "$numa1_cpus" "$rpc_pattern"
            done
        else
            log_warning "未找到进程: $rpc_pattern"
        fi
    done

    #绑定 collector进程
    local collector_pid=$(find_npu_collector_pid_with_top)
    
    bind_process_and_threads "$collector_pid" "$numa1_cpus" "collector" 

    local node_pid=$(find_node_pid_with_top)
    bind_process_and_threads "$node_pid" "$numa1_cpus" "node"

    # 绑定Main-Ark进程
    local ark_pids=$(pgrep -f "xLLM:Main-Ark")
    if [[ -n "$ark_pids" ]]; then
        echo "$ark_pids" | while read -r pid; do
            bind_process "$pid" "$numa1_cpus" "xLLM:Main-Ark"
        done
    else
        log_warning "未找到xLLM:Main-Ark进程"
    fi
}

# 显示当前绑定状态
show_binding_status() {
    log_info "显示当前绑定状态..."
    
    echo -e "\n${BLUE}=== CPU绑定状态摘要 ===${NC}"
    
    # 检查worker进程
    echo -e "\n${YELLOW}Worker进程状态:${NC}"
    for i in {0..15}; do
        local process_pattern="xLLM::XWorker_DP_$i"
        local pids=$(pgrep -f "xLLM::XWorker_DP_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM::XWorker_DP_${i}$")
        if [[ -n "$pids" ]]; then
            echo "$pids" | while read -r pid; do
                local current_affinity=$(taskset -cp "$pid" 2>/dev/null | awk '{print $6}')
                echo "  $process_pattern (PID: $pid) -> CPU: $current_affinity"
            done
        fi
    done
    
    # 检查driver进程
    echo -e "\n${YELLOW}Driver进程状态:${NC}"
    for i in {0..15}; do
        local process_pattern="xLLM::Driver_$i"
        local pids=$(pgrep -f "xLLM::Driver_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM::Driver_${i}$")
        if [[ -n "$pids" ]]; then
            echo "$pids" | while read -r pid; do
                local current_affinity=$(taskset -cp "$pid" 2>/dev/null | awk '{print $6}')
                echo "  $process_pattern (PID: $pid) -> CPU: $current_affinity"
            done
        fi
    done
    
    # 检查其他进程
    echo -e "\n${YELLOW}其他进程状态:${NC}"
    for pattern in "xLLM::Servicer" "xLLM:Main-Ark"; do
        local pids=$(pgrep -f "$pattern")
        if [[ -n "$pids" ]]; then
            echo "$pids" | while read -r pid; do
                local current_affinity=$(taskset -cp "$pid" 2>/dev/null | awk '{print $6}')
                echo "  $pattern (PID: $pid) -> CPU: $current_affinity"
            done
        fi
    done
    
    for i in {0..15}; do
        local rpc_pattern="xLLM:RpcProxy_$i"
        local pids=$(pgrep -f "xLLM:RpcProxy_${i}[^0-9]" 2>/dev/null || pgrep -f "xLLM:RpcProxy_${i}$")
        if [[ -n "$pids" ]]; then
            echo "$pids" | while read -r pid; do
                local current_affinity=$(taskset -cp "$pid" 2>/dev/null | awk '{print $6}')
                echo "  $rpc_pattern (PID: $pid) -> CPU: $current_affinity"
            done
        fi
    done

    local collector_pid=$(find_npu_collector_pid_with_top)
    local current_collector_affinity=$(taskset -cp "$collector_pid" 2>/dev/null | awk '{print $6}')
    echo "  collector (PID: $collector_pid) -> CPU: $current_collector_affinity"
}

# 主函数
main() {
    echo -e "${BLUE}=== xLLM进程CPU绑核脚本 ===${NC}"
    echo -e "机器配置: 256核，8个NUMA，每个NUMA 40核"
    echo -e "使用范围: NUMA 0-5 (前6个NUMA)"
    echo -e "绑定策略:"
    echo -e "  - NUMA 0: 中断处理"
    echo -e "  - CPU(40-59): 其他进程 (Servicer, RpcProxy, Main-Ark)"
    echo -e "  - CPU(60-123): Worker进程"
    echo -e "  - CPU(124-187): Driver进程"
    echo -e "  - CPU(192-252): Forward线程"
    echo ""
    
    # 检查环境
    check_root
    check_dependencies
    
    # 显示当前CPU信息
    log_info "当前系统CPU信息:"
    echo "  CPU总数: $(nproc)"
    echo "  NUMA节点数: $(lscpu | grep "NUMA node(s)" | awk '{print $3}')"
    echo ""
    
    # 开始绑定进程
    log_info "开始执行CPU绑核操作..."

    # 首先绑定中断
    bind_interrupts_to_numa0
    echo ""
    
    bind_worker_processes
    echo ""
    
    bind_driver_processes  
    echo ""
    
    bind_other_processes
    echo ""

    show_binding_status
    
    # 显示绑定结果
    show_binding_status
    
    log_success "CPU绑核操作完成！"
}

# 脚本入口点
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi
